#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脑电控制机器狗模型
基于SalientSleepNet架构的轻量化模型
"""

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import yaml

class EEGControlModel(keras.Model):
    """
    脑电控制模型
    基于SalientSleepNet架构，针对控制任务优化
    """
    
    def __init__(self, **kwargs):
        super(EEGControlModel, self).__init__()
        
        # 模型参数
        self.sequence_length = kwargs.get('sequence_length', 5)  # 5秒
        self.time_points = kwargs.get('time_points', 500)       # 500个时间点
        self.num_classes = 4  # 4个控制类别
        
        # 控制类别
        self.control_classes = {
            0: "FORWARD",    # 向前
            1: "BACKWARD",   # 向后  
            2: "LIE_DOWN",   # 趴下
            3: "TURN_AROUND" # 转圈
        }
        
        # 构建模型
        self.build_model()
    
    def build_model(self):
        """
        构建模型架构
        """
        # 输入层
        self.input_layer = layers.Input(shape=(self.sequence_length, self.time_points, 1, 1))
        
        # 特征提取层
        self.feature_extractor = self.build_feature_extractor()
        
        # 分类层
        self.classifier = self.build_classifier()
        
        # 输出层
        self.output_layer = layers.Dense(self.num_classes, activation='softmax')
    
    def build_feature_extractor(self):
        """
        构建特征提取器
        """
        model = keras.Sequential([
            # 第一层卷积 - 提取基本特征
            layers.Conv2D(16, (5, 1), padding='same', activation='relu'),
            layers.BatchNormalization(),
            layers.MaxPooling2D((2, 1)),
            
            # 第二层卷积 - 提取高级特征
            layers.Conv2D(32, (5, 1), padding='same', activation='relu'),
            layers.BatchNormalization(),
            layers.MaxPooling2D((2, 1)),
            
            # 第三层卷积 - 提取抽象特征
            layers.Conv2D(64, (5, 1), padding='same', activation='relu'),
            layers.BatchNormalization(),
            layers.MaxPooling2D((2, 1)),
            
            # 注意力机制
            layers.GlobalAveragePooling2D(),
            layers.Dense(64 // 4, activation='relu'),
            layers.Dense(64, activation='sigmoid'),
            layers.Reshape((1, 1, 64)),
            layers.Multiply()
        ])
        
        return model
    
    def build_classifier(self):
        """
        构建分类器
        """
        model = keras.Sequential([
            layers.Dense(128, activation='relu'),
            layers.Dropout(0.3),
            layers.Dense(64, activation='relu'),
            layers.Dropout(0.3),
            layers.Dense(32, activation='relu'),
            layers.Dense(self.num_classes, activation='softmax')
        ])
        
        return model
    
    def call(self, inputs):
        """
        前向传播
        """
        # 特征提取
        features = self.feature_extractor(inputs)
        
        # 分类
        output = self.classifier(features)
        
        return output

class LightweightEEGControlModel(keras.Model):
    """
    轻量化脑电控制模型
    适用于实时控制
    """
    
    def __init__(self, **kwargs):
        super(LightweightEEGControlModel, self).__init__()
        
        self.sequence_length = kwargs.get('sequence_length', 5)  # 5秒
        self.time_points = kwargs.get('time_points', 500)       # 500个时间点
        self.num_classes = 4
        
        # 控制类别
        self.control_classes = {
            0: "FORWARD",    # 向前
            1: "BACKWARD",   # 向后  
            2: "LIE_DOWN",   # 趴下
            3: "TURN_AROUND" # 转圈
        }
        
        self.build_model()
    
    def build_model(self):
        """
        构建轻量化模型
        """
        # 输入层
        self.input_layer = layers.Input(shape=(self.sequence_length, self.time_points, 1, 1))
        
        # 简化的特征提取
        self.feature_extractor = keras.Sequential([
            layers.Conv2D(8, (3, 1), padding='same', activation='relu'),
            layers.BatchNormalization(),
            layers.MaxPooling2D((2, 1)),
            
            layers.Conv2D(16, (3, 1), padding='same', activation='relu'),
            layers.BatchNormalization(),
            layers.MaxPooling2D((2, 1)),
            
            layers.GlobalAveragePooling2D()
        ])
        
        # 简化的分类器
        self.classifier = keras.Sequential([
            layers.Dense(32, activation='relu'),
            layers.Dropout(0.2),
            layers.Dense(16, activation='relu'),
            layers.Dense(self.num_classes, activation='softmax')
        ])
    
    def call(self, inputs):
        features = self.feature_extractor(inputs)
        output = self.classifier(features)
        return output

class RealTimeEEGController:
    """
    实时脑电控制器
    """
    
    def __init__(self, model_path, dog_controller):
        self.model = self.load_model(model_path)
        self.dog_controller = dog_controller
        self.buffer = []
        self.buffer_size = 500  # 5秒数据
        self.control_classes = {
            0: "FORWARD",
            1: "BACKWARD", 
            2: "LIE_DOWN",
            3: "TURN_AROUND"
        }
        
    def load_model(self, model_path):
        """
        加载训练好的模型
        """
        try:
            model = tf.keras.models.load_model(model_path)
            print(f"✅ 模型加载成功: {model_path}")
            return model
        except Exception as e:
            print(f"❌ 模型加载失败: {e}")
            return None
    
    def process_realtime_eeg(self, eeg_data):
        """
        处理实时脑电数据
        """
        # 添加到缓冲区
        self.buffer.append(eeg_data)
        
        # 保持缓冲区大小
        if len(self.buffer) > self.buffer_size:
            self.buffer.pop(0)
        
        # 当缓冲区满时进行预测
        if len(self.buffer) == self.buffer_size:
            prediction = self.predict_control()
            if prediction is not None:
                self.execute_control(prediction)
    
    def predict_control(self):
        """
        预测控制指令
        """
        if self.model is None:
            print("❌ 模型未加载")
            return None
        
        try:
            # 准备输入数据
            input_data = np.array(self.buffer).reshape(1, 5, 500, 1, 1)
            
            # 模型预测
            prediction = self.model.predict(input_data, verbose=0)
            control_class = np.argmax(prediction[0])
            confidence = np.max(prediction[0])
            
            # 只有当置信度超过阈值时才执行控制
            if confidence > 0.7:
                return control_class
            else:
                print(f"置信度太低: {confidence:.3f}")
                return None
                
        except Exception as e:
            print(f"❌ 预测失败: {e}")
            return None
    
    def execute_control(self, control_class):
        """
        执行控制指令
        """
        control_actions = {
            0: self.dog_controller.move_forward,
            1: self.dog_controller.move_backward,
            2: self.dog_controller.lie_down,
            3: self.dog_controller.turn_around
        }
        
        # 执行对应动作
        if control_class in control_actions:
            action_name = ['向前', '向后', '趴下', '转圈'][control_class]
            print(f"🎯 执行控制: {action_name}")
            control_actions[control_class]()
        else:
            print(f"❌ 未知控制类别: {control_class}")

class DogController:
    """
    机器狗控制器
    """
    
    def __init__(self, dog_ip="192.168.1.100", dog_port=8080):
        self.dog_ip = dog_ip
        self.dog_port = dog_port
        self.connection = self.connect_to_dog()
        
    def connect_to_dog(self):
        """
        连接到机器狗
        """
        try:
            # 这里需要根据具体的机器狗API进行连接
            # 例如使用socket或HTTP请求
            print(f"🔗 连接到机器狗: {self.dog_ip}:{self.dog_port}")
            # 模拟连接成功
            return True
        except Exception as e:
            print(f"❌ 连接失败: {e}")
            return False
    
    def move_forward(self):
        """
        向前移动
        """
        if self.connection:
            command = {"action": "forward", "duration": 2.0}
            self.send_command(command)
            print("🐕 机器狗向前移动")
        else:
            print("❌ 未连接到机器狗")
    
    def move_backward(self):
        """
        向后移动
        """
        if self.connection:
            command = {"action": "backward", "duration": 2.0}
            self.send_command(command)
            print("🐕 机器狗向后移动")
        else:
            print("❌ 未连接到机器狗")
    
    def lie_down(self):
        """
        趴下
        """
        if self.connection:
            command = {"action": "lie_down", "duration": 3.0}
            self.send_command(command)
            print("🐕 机器狗趴下")
        else:
            print("❌ 未连接到机器狗")
    
    def turn_around(self):
        """
        转圈
        """
        if self.connection:
            command = {"action": "turn_around", "duration": 3.0}
            self.send_command(command)
            print("🐕 机器狗转圈")
        else:
            print("❌ 未连接到机器狗")
    
    def send_command(self, command):
        """
        发送控制指令
        """
        # 这里实现具体的通信协议
        # 例如HTTP POST请求或TCP socket
        print(f"📡 发送指令: {command}")
        
        # 模拟发送指令
        import time
        time.sleep(0.1)  # 模拟网络延迟

def create_model(config):
    """
    根据配置创建模型
    """
    model_type = config.get('model', {}).get('type', 'lightweight')
    
    if model_type == 'lightweight':
        model = LightweightEEGControlModel(
            sequence_length=config['model']['sequence_length'],
            time_points=config['model']['time_points'],
            num_classes=config['model']['num_classes']
        )
    else:
        model = EEGControlModel(
            sequence_length=config['model']['sequence_length'],
            time_points=config['model']['time_points'],
            num_classes=config['model']['num_classes']
        )
    
    return model

def compile_model(model, config):
    """
    编译模型
    """
    optimizer = tf.keras.optimizers.Adam(
        learning_rate=config['training']['learning_rate']
    )
    
    loss_function = 'categorical_crossentropy'
    
    model.compile(
        optimizer=optimizer,
        loss=loss_function,
        metrics=['accuracy']
    )
    
    return model

def main():
    """
    主函数 - 演示模型创建
    """
    print("🧠 脑电控制机器狗模型")
    print("=" * 50)
    
    # 配置参数
    config = {
        'model': {
            'type': 'lightweight',
            'sequence_length': 5,
            'time_points': 500,
            'num_classes': 4
        },
        'training': {
            'epochs': 100,
            'batch_size': 16,
            'learning_rate': 0.001,
            'optimizer': 'adam'
        }
    }
    
    # 创建模型
    print("🏗️ 创建轻量化模型...")
    model = create_model(config)
    
    # 编译模型
    print("🔧 编译模型...")
    model = compile_model(model, config)
    
    # 打印模型结构
    print("📊 模型结构:")
    model.summary()
    
    # 创建模拟数据
    print("📈 创建模拟数据...")
    sample_data = np.random.randn(1, 5, 500, 1, 1)
    
    # 测试模型
    print("🧪 测试模型...")
    prediction = model.predict(sample_data, verbose=0)
    control_class = np.argmax(prediction[0])
    confidence = np.max(prediction[0])
    
    print(f"预测结果: 类别 {control_class}, 置信度 {confidence:.3f}")
    
    # 创建机器狗控制器
    print("🐕 创建机器狗控制器...")
    dog_controller = DogController()
    
    # 创建实时控制器
    print("🎮 创建实时控制器...")
    eeg_controller = RealTimeEEGController(
        model_path=None,  # 这里应该是实际的模型路径
        dog_controller=dog_controller
    )
    
    print("✅ 模型创建完成！")

if __name__ == "__main__":
    main() 